package org.obm.cassandra.client;

import static me.prettyprint.cassandra.utils.StringUtils.bytes;
import static me.prettyprint.cassandra.utils.StringUtils.string;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import me.prettyprint.cassandra.service.CassandraClient;
import me.prettyprint.cassandra.service.CassandraClientPool;
import me.prettyprint.cassandra.service.Keyspace;

import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.ColumnPath;
import org.apache.cassandra.thrift.InvalidRequestException;
import org.apache.cassandra.thrift.NotFoundException;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.cassandra.thrift.TimedOutException;
import org.apache.cassandra.thrift.UnavailableException;
import org.apache.thrift.TException;

/**
 * Helper functions for Cassandra data manipulation
 * 
 * @author tom
 * 
 */
public final class CassHelper {

	public static final void cleanup(CassandraClient cli) {
		if (cli != null) {
			try {
				Activator.getDefault().getPool().releaseClient(cli);
			} catch (Exception e) {
			}
		}
	}

	/**
	 * Inserts :
	 * 
	 * <pre>
	 * columnFamily {
	 * 		key {
	 * 			mapKey1 : data.get(mapKey1)
	 * 			mapKey2 : data.get(mapKey2)
	 *  			...
	 * 			mapKeyN : data.get(mapKeyN)
	 * 		}
	 * }
	 * </pre>
	 * 
	 * @param ks
	 *            the Key space
	 * @param columnFamily
	 *            the column family to use
	 * @param key
	 *            the new key to insert
	 * @param data
	 * @throws InvalidRequestException
	 * @throws UnavailableException
	 * @throws TException
	 * @throws TimedOutException
	 */
	public static final void insert(Keyspace ks, String columnFamily,
			String key, Map<String, String> data)
			throws InvalidRequestException, UnavailableException, TException,
			TimedOutException {
		for (String s : data.keySet()) {
			ColumnPath cp = cp(columnFamily, s);
			ks.insert(key, cp, bytes(data.get(s)));
		}
	}

	public static final ColumnPath cp(String cf, String key) {
		ColumnPath ret = new ColumnPath(cf);
		ret.setColumn(bytes(key));
		return ret;
	}

	/**
	 * Fetches the columns in key as a map :
	 * 
	 * <pre>
	 * family {
	 * 		key {
	 * 			mapKey1 : data.get(mapKey1)
	 * 			mapKey2 : data.get(mapKey2)
	 *  			...
	 * 			mapKeyN : data.get(mapKeyN)
	 * 		}
	 * }
	 * </pre>
	 * 
	 * @param ks
	 *            the Key space
	 * @param family
	 *            the column family
	 * @param key
	 * @return all columns name & values in a map
	 * @throws InvalidRequestException
	 * @throws NotFoundException
	 * @throws UnavailableException
	 * @throws TException
	 * @throws TimedOutException
	 * 
	 */
	public static final Map<String, String> fetchAll(Keyspace ks,
			String family, String key) throws InvalidRequestException,
			NotFoundException, UnavailableException, TException,
			TimedOutException {

		ColumnParent clp = new ColumnParent(family);
		SliceRange sr = new SliceRange(new byte[0], new byte[0], false, 150);
		SlicePredicate sp = new SlicePredicate();
		sp.setSlice_range(sr);

		List<Column> values = ks.getSlice(key, clp, sp);
		HashMap<String, String> ret = new HashMap<String, String>(2 * values
				.size());
		for (Column c : values) {
			ret.put(string(c.getName()), string(c.getValue()));
		}
		return ret;
	}

	public static final boolean run(String ksName, CassCommand cmd) {
		CassandraClientPool pool = org.obm.cassandra.client.Activator
				.getDefault().getPool();

		CassandraClient client = null;
		try {
			client = pool.borrowClient();
			Keyspace ks = client.getKeyspace(ksName);
			cmd.run(ks);
			return true;
		} catch (Throwable t) {
			return false;
		} finally {
			CassHelper.cleanup(client);
		}
	}

}
